import math
import numpy as np
from collections import Counter
from sklearn.metrics.pairwise import cosine_similarity
from nltk.stem import PorterStemmer
from sklearn.feature_extraction.text import TfidfVectorizer


# =====================BLEU SCORE=====================
class BLEU:

    def __init__(self, n_grams=(1, 2, 3, 4)):
        """
        初始化BLEU评分类。
        :param n_grams: n元语法的元组，默认为1到4。
        """
        self.n_grams = n_grams

    def _ngram_precision(self, reference, translation, n):
        """
        计算n元语法的精确度。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :param n: n元语法中的n。
        :return: n元语法的精确度。
        """
        ref_ngrams = Counter([tuple(reference[i:i + n]) for i in range(len(reference) - n + 1)])
        hyp_ngrams = Counter([tuple(translation[i:i + n]) for i in range(len(translation) - n + 1)])
        match_count = sum((hyp_ngrams & ref_ngrams).values())
        total_count = sum(hyp_ngrams.values())
        return match_count / total_count if total_count > 0 else 0

    def _brevity_penalty(self, reference, translation):
        """
        计算简洁性惩罚。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: 简洁性惩罚值。
        """
        ref_len, hyp_len = len(reference), len(translation)
        return 1 if hyp_len > ref_len else math.exp(1 - ref_len / hyp_len) if hyp_len > 0 else 0

    def compute(self, reference, translation):
        """
        计算BLEU评分。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: BLEU评分。
        """
        precisions = [self._ngram_precision(reference, translation, n) for n in self.n_grams]
        geo_mean = math.exp(sum(math.log(p) for p in precisions if p > 0) / len(precisions)) if precisions else 0
        return self._brevity_penalty(reference, translation) * geo_mean


# ========================METEOR======================
class METEOR:
    """
    METEOR评分类。
    用于计算翻译文本相对于参考文本的METEOR评分。
    """
    def __init__(self):
        self.stemmer = PorterStemmer()

    def _exact_match(self, reference, translation):
        """
        计算精确匹配的数量。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: 精确匹配的数量。
        """
        # 按词频计算精确匹配数量
        ref_counts = Counter(reference)
        hpy_counts = Counter(translation)
        return sum(min(ref_counts[word], hpy_counts[word]) for word in ref_counts.keys())

    def _stem_match(self, reference, translation, exact_match_count):
        """
        计算词干匹配的数量。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :param exact_match_count: 精确匹配的数量。
        :return: 词干匹配的数量。
        """
        # 按词频计算词干匹配数量。排除已精确匹配的词
        ref_stems = Counter(self.stemmer.stem(w) for w in reference)
        hyp_stems = Counter(self.stemmer.stem(w) for w in translation)
        # 计算匹配数量，但要排除已经完全匹配的单词
        stem_match_count = sum(min(ref_stems[stem], hyp_stems[stem]) for stem in ref_stems.keys())
        # 避免重复计数
        return max(0, stem_match_count - exact_match_count)

    def compute(self, reference, translation):
        """
        计算METEOR评分。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: METEOR评分。
        """
        exact_match_count = self._exact_match(reference, translation)
        stem_match_count = self._stem_match(reference, translation, exact_match_count)
        matches = exact_match_count + stem_match_count
        precision = matches / len(translation) if translation else 0
        recall = matches / len(reference) if reference else 0
        if precision == 1.0 and recall == 1.0:
            return 1.0
        f_mean = (10 * precision * recall) / (recall + 9 * precision) if (precision + recall) > 0 else 0
        return f_mean


# =======================TER==================
class TER:
    """
    TER（Translation Edit Rate）评分类。
    用于计算翻译文本相对于参考文本的编辑距离。
    """
    def _edit_distance(self, ref, hyp):
        """
        计算编辑距离。
        :param ref: 参考文本，词列表。
        :param hyp: 翻译文本，词列表。
        :return: 编辑距离。
        """
        dp = [[0] * (len(hyp) + 1) for _ in range(len(ref) + 1)]
        for i in range(len(ref) + 1):
            for j in range(len(hyp) + 1):
                if i == 0:
                    dp[i][j] = j
                elif j == 0:
                    dp[i][j] = i
                else:
                    cost = 0 if ref[i - 1] == hyp[j - 1] else 1
                    dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j - 1] + 1, dp[i - 1][j - 1] + cost)
        return dp[-1][-1]

    def compute(self, reference, translation):
        """
        计算TER评分。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: TER评分。
        """
        return self._edit_distance(reference, translation) / len(reference)


# ===================== chrF ======================
class ChrF:
    """
    chrF评分类。
    用于计算翻译文本相对于参考文本的字符n元语法的F分数。
    """
    def __init__(self, beta=2):
        """
        初始化chrF评分类。
        :param beta: F分数中的beta参数，默认为2。
        """
        self.beta = beta

    def _character_ngrams(self, sentence, n):
        """
        生成字符n元语法列表。
        :param sentence: 输入的句子，词列表。
        :param n: n元语法中的n。
        :return: 字符n元语法列表。
        """
        sentence = "".join(sentence)  # 将词列表拼接成字符串
        return [sentence[i:i + n] for i in range(len(sentence) - n + 1)]

    def _chrf_precision_recall(self, reference, translation, n):
        """
         计算字符n元语法的精确度和召回率。
         :param reference: 参考文本，词列表。
         :param translation: 翻译文本，词列表。
         :param n: n元语法中的n。
         :return: 字符n元语法的精确度和召回率。
         """
        ref_ngrams = set(self._character_ngrams(reference, n))
        hyp_ngrams = set(self._character_ngrams(translation, n))
        match_count = len(ref_ngrams & hyp_ngrams)
        precision = match_count / len(hyp_ngrams) if hyp_ngrams else 0
        recall = match_count / len(ref_ngrams) if ref_ngrams else 0
        return precision, recall

    def compute(self, reference, translation):
        """
        计算chrF评分。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: chrF评分。
        """
        precisions, recalls = zip(*[self._chrf_precision_recall(reference, translation, n) for n in range(1, 7)])
        precision, recall = sum(precisions) / 6, sum(recalls) / 6
        return (1 + self.beta ** 2) * (precision * recall) / (
                self.beta ** 2 * precision + recall) if precision and recall else 0


# ====================== COMET ====================
class COMET:
    """
    COMET评分类。
    用于计算翻译文本相对于参考文本的COMET评分，基于TF-IDF和余弦相似度。
    """
    def __init__(self):
        self.vectorizer = TfidfVectorizer()

    def compute(self, reference, translation):
        """
        计算COMET评分。
        :param reference: 参考文本，词列表。
        :param translation: 翻译文本，词列表。
        :return: COMET评分。
        """
        ref_text = " ".join(reference)
        hyp_text = " ".join(translation)

        # 计算TF-IDF矩阵
        tfidf_matrix = self.vectorizer.fit_transform([ref_text, hyp_text])

        # 计算余弦相似度
        return cosine_similarity(tfidf_matrix[0], tfidf_matrix[1])[0][0]


# TODO: 把测试代码移到测试文件中
# ========================= 测试 ===========================
if __name__ == '__main__':
    ref = "the cat is on the mat".split()
    hyp = "the cat is on mat".split()

    bleu = BLEU().compute(ref, hyp)
    print(f"BLEU: {bleu:.4f}")
    meteor = METEOR().compute(ref, hyp)
    print(f"METEOR: {meteor:.4f}")
    ter = TER().compute(ref, hyp)
    print(f"TER: {ter:.4f}")
    chrf = ChrF().compute(ref, hyp)
    print(f"chrF: {chrf:.4f}")
